import os
import sys
import csv
import subprocess
import copy
import datetime
 
def read_csv(data_path):
    with open(data_path, "r") as f:
        data = csv.reader(f)
        list_data = list(data)
    return list_data
 
 
class OpProfiling():
    def __init__(self):
        self.no_list = ["Model ID", "Stream ID", "Task Start Time", "Input Shapes", "Input Data Types", "Input Formats", \
                        "Output Shapes", "Output Data Types", "Output Formats", "total_cycles", "mac_time(us)", "scalar_time(us)", \
                        "mte1_time(us)", "mte2_time(us)", "mte3_time(us)", "icache_miss_rate", "memory_bound"]
        self.yes_list = []

    def get_profiling_info(self, csv_data):
        param_num = len(csv_data[0])
        task_num = len(csv_data)
        param = []
        values = []
        for i in range(param_num):
            if self.yes_list != []:
                if str(csv_data[0][i]) in self.yes_list:
                    param.append(str(csv_data[0][i]))
                continue
            if str(csv_data[0][i]) not in self.no_list:
                param.append(str(csv_data[0][i]))
        for q in range(1, task_num):
            for z in range(param_num):
                if "_ratio" in csv_data[0][z] and float(csv_data[q][z]) > 0.8:
                    print("\033[93m[> 警告 <]\033[0m", "性能结果中,存在ratio较高的情况", "\033[93m%s:%.3f\033[0m"%(csv_data[0][z], float(csv_data[q][z])))
                if self.yes_list != []:
                    if str(csv_data[0][z]) in self.yes_list:
                        values.append(str(csv_data[q][z]))
                    continue
                if str(csv_data[0][z]) not in self.no_list:
                    values.append(str(csv_data[q][z]))
        return param, values
 
    def get_length(self, header_names, data):
        size = []
        for i in range(len(header_names)):
            col_data = [len(header_names[i])]
            for d in data:
                col_data.append(len(str(d)))
            size.append(max(col_data) + 4)
        return max(size)
 
    def format_str(self, info, length):
        info_strs = []
        for i in range(len(info)):
            prof_data = info[i]
            info_str = "  " + str(prof_data) + " " * (length - len(str(prof_data)) - 2)
            info_strs.append(info_str)
        res = "|" + "|".join(info_strs) + "|"
        return res
 
    def format_print(self, prof_res, values, one_line_num=6):
        header_name = prof_res
        data = values
        task_num = len(data) // len(header_name)
        param_num = len(header_name)
        lines_num = (param_num + one_line_num - 1) // one_line_num
        length = self.get_length(header_name, data)
        
        for i in range(lines_num):
            format_strs = []
            if i == lines_num - 1:
                line_name = header_name[i * one_line_num :]
                line_data = data[i * one_line_num :]
                format_strs.append(self.format_str(line_name, length))
                for j in range(task_num):
                    start_index = i * one_line_num + j * param_num 
                    line_data = data[start_index : start_index + len(line_name)]
                    format_strs.append(self.format_str(line_data, length))

                base_line = ["-" * length for i in range(len(line_name))]
                start_end_line = str(i) + str(i).join(base_line) + str(i)
                print(start_end_line)
                for i in format_strs:
                    print(i)
                print(start_end_line)
            else:
                line_name = header_name[i * one_line_num : i * one_line_num + one_line_num]
                format_strs.append(self.format_str(line_name, length))
                for j in range(task_num):
                    start_index = i * one_line_num + j * param_num 
                    line_data = data[start_index : start_index + one_line_num]
                    format_strs.append(self.format_str(line_data, length))

                base_line = ["-" * length for i in range(one_line_num)]
                start_end_line = str(i) + str(i).join(base_line) + str(i)
                print(start_end_line)
                for i in format_strs:
                    print(i)
                print(start_end_line)
 
    def op_profiling(self, data_path, device_id=0):
        ms_path = os.environ.get("ASCEND_TOOLKIT_HOME", os.environ.get("ASCEND_HOME_PATH", ""))
        ms_py = "tools/profiler/profiler_tool/analysis/msprof/msprof.py"
        print_py = "tools/profiler/profiler_tool/analysis/common_func/common.py"
        print_path = os.path.join(ms_path, print_py)
        with open(print_path, "a+") as f:
            f.writelines("\nsys.stdout = open(os.devnull, 'w')")
            # f.writelines("\nsys.stdout = sys.__stdout__")
        names = os.listdir(data_path)
        job_path = None
        res, values = [], []
        for name in names:
            if "PROF" in name or "JOB" in name:
                job_path = os.path.join(data_path, name)
                break
 
        if job_path:
            msprof_cmd = "python3 %s export summary -dir=%s" % (os.path.join(ms_path, ms_py), job_path)
            subprocess.run(msprof_cmd, shell=True)
            msprof_timeline_cmd = "python3 %s export timeline -dir=%s" % (os.path.join(ms_path, ms_py), job_path)
            subprocess.run(msprof_timeline_cmd, shell=True)
            summary_path = os.path.join(job_path, "device_%s/summary" % device_id)
            if not os.path.exists(summary_path):
                summary_path = os.path.join(job_path, "mindstudio_profiler_output/")
            if os.path.exists(summary_path):
                names = os.listdir(summary_path)
                result_name = None
                for name in names:
                    if "op_summary" in name:
                        result_name = name
                        break
 
                if not result_name:
                    print("\033[31m[> 错误 <]\033[0m" + "性能收集没有op_summary ", names)
                    print("\033[31m[> 错误 <]\033[0m" + "1.当前device_id=%d,请尝试换卡"%device_id)
                    return res, values
                csv_data = read_csv(os.path.join(summary_path, result_name))
                res, values = self.get_profiling_info(csv_data)
            else:
                print("\033[31m[> 错误 <]\033[0m" + "性能数据输出没有mindstudio_profiler_output")
                print("\033[31m[> 错误 <]\033[0m" + "1.当前device_id=%d,请尝试换卡"%device_id)
            msprof_cmd = "mv %s %s/prof_%s" %(job_path, data_path, datetime.datetime.now().strftime("%M_%S"))
            subprocess.run(msprof_cmd, shell=True)
        else:
            print("\033[31m[> 错误 <]\033[0m" + "没有性能目录:", job_path)
            print("\033[31m[> 错误 <]\033[0m" + "1.当前device_id=%d,请尝试换卡"%device_id)
        return res, values
 
 
op_profiling = OpProfiling()